import os

import torch
from natsort import natsorted

import onnxruntime
import torchvision.transforms as transforms
from PIL import Image


def preprocess_image(image):
    image = image.convert('RGB')
    transform = transforms.Compose([
        transforms.Resize((448, 448)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    return transform(image).unsqueeze(0)


# Define inference function using onnxruntime
def predict(ort_session, image, classlist):
    image = preprocess_image(image)
    ort_inputs = {ort_session.get_inputs()[0].name: image.numpy()}
    ort_outs = ort_session.run(None, ort_inputs)
    pred = torch.from_numpy(ort_outs[0]).squeeze()
    id = pred.argmax()
    # shipclass = ['none', 'SH_Q', 'SH_Z', 'JZX_Z', 'YL_Q', 'YL_Z', 'ZF', 'GC']
    name = classlist[id]
    conf = pred[id]
    return name, conf


""""
测试模型效果，也可作为使用的基础
"""

if __name__ == '__main__':

    """"
    测试的参数与训练保持一致
    """
    input_shape = (1, 3, 448, 448)
    onnx_path = 'pth/best_model.onnx'
    ort_session = onnxruntime.InferenceSession(onnx_path)

    """"
    folder_path: 测试文件夹，测试文件夹下所有图片，输出类型
    """
    folder_path = "test"
    files = []
    lists = os.listdir(folder_path)
    classlist = ['none', 'SH_Q', 'SH_Z', 'JZX_Z', 'YL_Q', 'YL_Z', 'ZF', 'GC']
    # 按文件名中的数字排序
    sorted_files = natsorted(lists)
    for filename in sorted_files:
        if filename.endswith(".jpg") or filename.endswith(".png"):
            img = Image.open(os.path.join(folder_path, filename))
            name, conf = predict(ort_session, img, classlist)
            print(filename, "-", name, conf)


